/// -----------------------------------------------------------------------
/// (c) 2009 Microsoft Corporation -- All rights reserved
/// This code is licensed under the MS-PL
/// http://www.opensource.org/licenses/ms-pl.html
/// Courtesy of the Open Source Techology Center: http://port25.technet.com
/// -----------------------------------------------------------------------

#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <string.h>
#include <detours.h>
#include "trace.h"

#if (_MSC_VER < 1299)
typedef ULONG * PULONG_PTR;
typedef ULONG ULONG_PTR;
typedef LONG * PLONG_PTR;
typedef LONG LONG_PTR;
#endif

//////////////////////////////////////////////////////////////////////////////
#pragma warning(disable:4127)   // Many of our asserts are constants.

#define ASSERT_ALWAYS(x)   \
    do {                                                        \
    if (!(x)) {                                                 \
            AssertMessage(#x, __FILE__, __LINE__);              \
            DebugBreak();                                       \
    }                                                           \
    } while (0)

#ifndef NDEBUG
#define ASSERT(x)           ASSERT_ALWAYS(x)
#else
#define ASSERT(x)
#endif

#define UNUSED(c)       (c) = (c)
#define ARRAYOF(x)      (sizeof(x)/sizeof(x[0]))

//////////////////////////////////////////////////////////////////////////////
enum {
    CLIENT_AWAITING_PIPE_ACCEPT = 0x21,
    CLIENT_AWAITING_PIPE_DATA   = 0x22,
};

typedef struct _CLIENT : OVERLAPPED
{
    HANDLE          hPipe;
    LONG            nClient;
    HANDLE          hFile;
    BOOL            fAwaitingAccept;
    PVOID           Zero;
    TBLOG_MESSAGE   Message;

    BOOL LogMessage(PTBLOG_MESSAGE pMessage, DWORD nBytes);
    BOOL LogMessageV(PCHAR pszMsg, ...);
} CLIENT, *PCLIENT;

//////////////////////////////////////////////////////////////////////////////
//
CHAR        s_szLogFile[MAX_PATH];
CHAR*		s_szJustFileName;
CHAR        s_szTempPath[MAX_PATH];
CHAR        s_szPipe[MAX_PATH];
LONG        s_nActiveClients = 0;
LONG        s_nTotalClients = 0;
LONGLONG    s_llStartTime;
BOOL        s_fVerbose = FALSE;
TBLOG_PAYLOAD s_Payload;

//////////////////////////////////////////////////////////////////////////////
VOID MyErrExit(PCSTR pszMsg)
{
    DWORD error = GetLastError();

    fprintf(stderr, "TRACE: Error %d in %s.\n", error, pszMsg);
    fflush(stderr);
    exit(1);
}

//////////////////////////////////////////////////////////////////////////////
BOOL CLIENT::LogMessageV(PCHAR pszMsg, ...)
{
    DWORD cbWritten = 0;
    CHAR szBuf[1024];
    INT cbToWrite = 0;

    va_list args;
    va_start(args, pszMsg);
#ifdef _CRT_INSECURE_DEPRECATE
    cbToWrite += _vsnprintf_s(szBuf + cbToWrite,
                              sizeof(szBuf) - cbToWrite,
                              sizeof(szBuf) - cbToWrite - 1,
                              pszMsg, args);
#else
    cbToWrite += _vsnprintf(szBuf + cbToWrite, sizeof(szBuf) - cbToWrite, pszMsg, args);
#endif
    va_end(args);

#ifdef _CRT_INSECURE_DEPRECATE
    cbToWrite += _snprintf_s(szBuf + cbToWrite,
                             sizeof(szBuf) - cbToWrite,
                             sizeof(szBuf) - cbToWrite - 1,
                             "\n");
#else
    cbToWrite += _snprintf(szBuf + cbToWrite, sizeof(szBuf) - cbToWrite, "\n");
#endif

    if (cbToWrite < 0) {
        szBuf[sizeof(szBuf)-1] = '\n';
        cbToWrite = sizeof(szBuf);
    }

    WriteFile(hFile, szBuf, cbToWrite, &cbWritten, NULL);
    return TRUE;
}

BOOL CLIENT::LogMessage(PTBLOG_MESSAGE pMessage, DWORD nBytes)
{
    // Sanity check the size of the message.
    //
    if (nBytes > pMessage->nBytes) {
        nBytes = pMessage->nBytes;
    }
    if (nBytes >= sizeof(*pMessage)) {
        nBytes = sizeof(*pMessage) - 1;
    }

    // Don't log message if there isn't and message text.
    DWORD cbWrite = nBytes - offsetof(TBLOG_MESSAGE, szMessage);
    if (cbWrite <= 0 ) {
        return TRUE;
    }

    if (s_fVerbose) {
        printf("[%s]", pMessage->szMessage);
    }

    DWORD cbWritten = 0;
    WriteFile(hFile, pMessage->szMessage, cbWrite, &cbWritten, NULL);
    return TRUE;
}

BOOL CloseConnection(PCLIENT pClient)
{
    InterlockedDecrement(&s_nActiveClients);
    if (pClient != NULL) {
        if (pClient->hPipe != INVALID_HANDLE_VALUE) {
            //FlushFileBuffers(pClient->hPipe);
            if (!DisconnectNamedPipe(pClient->hPipe)) {
#if DEBUG
                DWORD error = GetLastError();
                pClient->LogMessageV("<!-- Error %d in DisconnectNamedPipe. -->\n", error);
#endif 
            }
            CloseHandle(pClient->hPipe);
            pClient->hPipe = INVALID_HANDLE_VALUE;
        }
        if (pClient->hFile != INVALID_HANDLE_VALUE) {
            CloseHandle(pClient->hFile);
            pClient->hFile = INVALID_HANDLE_VALUE;
        }
        GlobalFree(pClient);
        pClient = NULL;
    }
    return TRUE;
}

// Creates a pipe instance and initiate an accept request.
PCLIENT CreatePipeConnection(HANDLE hCompletionPort, LONG nClient)
{
    HANDLE hPipe = CreateNamedPipe(s_szPipe,                    // pipe name
                                   PIPE_ACCESS_INBOUND |        // read-only access
                                   FILE_FLAG_OVERLAPPED,        // overlapped mode
                                   PIPE_TYPE_MESSAGE |          // message-type pipe
                                   PIPE_READMODE_MESSAGE |      // message read mode
                                   PIPE_WAIT,                   // blocking mode
                                   PIPE_UNLIMITED_INSTANCES,    // unlimited instances
                                   0,                           // output buffer size
                                   0,                           // input buffer size
                                   20000,                       // client time-out
                                   NULL);                       // no security attributes
    if (hPipe == INVALID_HANDLE_VALUE) {
        MyErrExit("CreateNamedPipe");
    }

    // Allocate the client data structure.
    //
    PCLIENT pClient = (PCLIENT) GlobalAlloc(GPTR, sizeof(CLIENT));
    if (pClient == NULL) {
        MyErrExit("GlobalAlloc pClient");
    }

    CHAR szLogFile[MAX_PATH];
#ifdef _CRT_INSECURE_DEPRECATE
    sprintf_s(szLogFile, ARRAYOF(szLogFile), "%s\\%s.%08d.xml", s_szTempPath,  s_szJustFileName, nClient);
#else
    sprintf(szLogFile, "%s\\%s.%08d.xml", s_szTempPath, s_szJustFileName, nClient);
#endif

    ZeroMemory(pClient, sizeof(*pClient));
    pClient->hPipe = hPipe;
    pClient->nClient = nClient;
    pClient->fAwaitingAccept = TRUE;
    pClient->hFile = CreateFile(szLogFile,
                                GENERIC_WRITE,
                                FILE_SHARE_READ,
                                NULL,
                                CREATE_ALWAYS,
                                FILE_ATTRIBUTE_NORMAL |
                                FILE_FLAG_SEQUENTIAL_SCAN,
                                NULL);
    if (pClient->hFile == INVALID_HANDLE_VALUE) {
        fprintf(stderr, "TRACE: Error opening output file: %s: %d\n\n",
                szLogFile, GetLastError());
        fflush(stderr);
        MyErrExit("CreateFile");
    }

    // Associate file with our complietion port.
    if (!CreateIoCompletionPort(pClient->hPipe, hCompletionPort, (ULONG_PTR)pClient, 0)) {
        MyErrExit("CreateIoComplietionPort pClient");
    }

    if (!ConnectNamedPipe(hPipe, pClient)) {
        DWORD error = GetLastError();

        if (error == ERROR_IO_PENDING) {
            return NULL;
        }
        if (error == ERROR_PIPE_CONNECTED) {
#if 0
            pClient->LogMessageV("<!-- ConnectNamedPipe client already connected. -->");
#endif
            pClient->fAwaitingAccept = FALSE;
        }
        else if (error != ERROR_IO_PENDING &&
                 error != ERROR_PIPE_LISTENING) {

            MyErrExit("ConnectNamedPipe");
        }
    }
    else {
        fprintf(stderr, "*** ConnectNamedPipe accepted immediately.\n");
#if 0
        pClient->LogMessageV("<!-- ConnectNamedPipe accepted immediately. -->");
#endif
        pClient->fAwaitingAccept = FALSE;
    }
    return pClient;
}

BOOL DoRead(PCLIENT pClient)
{
    SetLastError(NO_ERROR);
    DWORD nBytes = 0;
    BOOL b = ReadFile(pClient->hPipe, &pClient->Message, sizeof(pClient->Message),
                      &nBytes, pClient);

    DWORD error = GetLastError();

    if (b && error == NO_ERROR) {
        return TRUE;
    }
    if (error == ERROR_BROKEN_PIPE) {
#if DEBUG
        pClient->LogMessageV("<!-- **** ReadFile 002 *** ERROR_BROKEN_PIPE [%d] -->\n", nBytes);
#endif
        CloseConnection(pClient);
        return TRUE;
    }
    else if (error == ERROR_INVALID_HANDLE) {
        // ?
#if DEBUG
        pClient->LogMessageV("<!-- **** ReadFile 002 *** ERROR_INVALID_HANDLE -->\n");
#endif 
        // I have no idea why this happens.  Our remedy is to drop the connection.
        return TRUE;
    }
    else if (error != ERROR_IO_PENDING) {
        if (b) {
#if DEBUG
            pClient->LogMessageV("<!-- **** ReadFile 002 succeeded: %d -->\n", error);
#endif
        }
        else {
#if DEBUG
            pClient->LogMessageV("<!-- **** ReadFile 002 failed: %d -->\n", error);
#endif
        }
        CloseConnection(pClient);
    }
    return TRUE;
}

DWORD WINAPI WorkerThread(LPVOID pvVoid)
{
    PCLIENT pClient;
    BOOL b;
    LPOVERLAPPED lpo;
    DWORD nBytes;
    HANDLE hCompletionPort = (HANDLE)pvVoid;

    for (BOOL fKeepLooping = TRUE; fKeepLooping;) {
        pClient = NULL;
        lpo = NULL;
        nBytes = 0;
        b = GetQueuedCompletionStatus(hCompletionPort,
                                      &nBytes, (PULONG_PTR)&pClient, &lpo, INFINITE);

        if (!b) {
            if (pClient) {
                if (GetLastError() == ERROR_BROKEN_PIPE) {
#if DEBUG
                    pClient->LogMessageV("<!-- Client closed pipe. -->");
#endif
                }
                else {
#if DEBUG
                    pClient->LogMessageV("<!-- *** GetQueuedCompletionStatus failed %d -->",GetLastError());
#endif
                }
                CloseConnection(pClient);
            }
            continue;
        }

        if (pClient->fAwaitingAccept) {
            BOOL fAgain = TRUE;
            while (fAgain) {
                LONG nClient = InterlockedIncrement(&s_nTotalClients);
                InterlockedIncrement(&s_nActiveClients);
                pClient->fAwaitingAccept = FALSE;

                PCLIENT pNew = CreatePipeConnection(hCompletionPort, nClient);

                fAgain = FALSE;
                if (pNew != NULL) {
                    fAgain = !pNew->fAwaitingAccept;
                    DoRead(pNew);
                }
            }
        }
        else {
            if (nBytes <= offsetof(TBLOG_MESSAGE, szMessage)) {
                pClient->LogMessageV("</Process>\n");
                CloseConnection(pClient);
                continue;
            }
            pClient->LogMessage(&pClient->Message, nBytes);
        }

        DoRead(pClient);
    }
    return 0;
}

BOOL CreateWorkers(HANDLE hCompletionPort)
{
    DWORD dwThread;
    HANDLE hThread;
    DWORD i;
    SYSTEM_INFO SystemInfo;

    GetSystemInfo(&SystemInfo);

    for (i = 0; i < 1; i++) {
        hThread = CreateThread(NULL, 0, WorkerThread, hCompletionPort, 0, &dwThread);
        if (!hThread) {
            MyErrExit("CreateThread WorkerThread");
            // Unreachable: return FALSE;
        }
        CloseHandle(hThread);
    }
    return TRUE;
}

DWORD CopyEnvironment(PWCHAR pwzzOut, PCWSTR pwzzIn)
{
    PCWSTR pwzzBeg = pwzzOut;
    while (*pwzzIn) {
        while (*pwzzIn) {
            *pwzzOut++ = *pwzzIn++;
        }
        *pwzzOut++ = *pwzzIn++;   // Copy zero.
    }
    *pwzzOut++ = '\0';    // Add last zero.

    return (DWORD)(pwzzOut - pwzzBeg);
}

void WriteOut( HANDLE hFileHandle , const char* text ) {
    DWORD sz = (DWORD) strlen(text);
    WriteFile(hFileHandle , text , sz , &sz, NULL);
}

//////////////////////////////////////////////////////////////////////////////
//
DWORD main(int argc, char **argv)
{
    HANDLE hCompletionPort;
    BOOL fNeedHelp = FALSE;
    WCHAR wzzDrop[1024] = L"build\0nmake\0";
	GetTempPath(MAX_PATH, s_szTempPath );
	s_szJustFileName = s_szLogFile;

	// give a decent default output file name.
#ifdef _CRT_INSECURE_DEPRECATE
    sprintf_s(s_szLogFile, ARRAYOF(s_szLogFile), "TraceLog[%d]", GetCurrentProcessId());
#else
	sprintf(s_szLogFile, "TraceLog[%d]", GetCurrentProcessId());
#endif

    GetSystemTimeAsFileTime((FILETIME *)&s_llStartTime);
#ifdef _CRT_INSECURE_DEPRECATE
    sprintf_s(s_szPipe, ARRAYOF(s_szPipe), "%s.%d", TBLOG_PIPE_NAME, GetCurrentProcessId());
#else
    sprintf(s_szPipe, "%s.%d", TBLOG_PIPE_NAME, GetCurrentProcessId());
#endif

    int arg = 1;
    for (; arg < argc && (argv[arg][0] == '-' || argv[arg][0] == '/'); arg++) {
        CHAR *argn = argv[arg] + 1;
        CHAR *argp = argn;
        while (*argp && *argp != ':') {
            argp++;
        }
        if (*argp == ':') {
            *argp++ = '\0';
        }

        switch (argn[0]) {

          case 'd':                                     // Drop Processes
          case 'D':
            if (*argp) {
                PWCHAR pwz = wzzDrop;
                while (*argp) {
                    if (*argp == ';') {
                        *pwz++ = '\0';
                    }
                    else {
                        *pwz++ = *argp++;
                    }
                }
                *pwz++ = '\0';
                *pwz = '\0';
            }
          case 'o':                                 // Output file.
          case 'O':
			if( *argp ) {
#ifdef _CRT_INSECURE_DEPRECATE
            strcpy_s(s_szLogFile, ARRAYOF(s_szLogFile), argp);
#else
            strcpy(s_szLogFile, argp);
#endif
			CHAR* posn = s_szJustFileName = s_szLogFile;
			while( *posn ) {
					if( *posn == '\\' && *(posn+1) != '\\' && *(posn+1) != 0x00 ) {
						s_szJustFileName  = posn+1;
						printf("%s",s_szJustFileName  );
					}
					posn++;
				}
			}

            break;

          case 'v':                                     // Verbose
          case 'V':
            s_fVerbose = TRUE;
            break;

          case '?':                                 // Help.
            fNeedHelp = TRUE;
            break;

          default:
            fNeedHelp = TRUE;
            printf("TRACE: Bad argument: %s:%s\n", argn, argp);
            break;
        }
    }

    if (arg >= argc) {
        fNeedHelp = TRUE;
    }

    if (fNeedHelp) {
        printf("Usage:\n"
               "    trace [options] command {command arguments}\n"
               "Options:\n"
               "    /o:file    Log all events to the output files.\n"
               "    /?         Display this help message.\n"
               "Summary:\n"
               "    Runs the build commands and figures out which files have dependencies..\n"
               "\n");
        exit(9001);
    }

    // Create the completion port.
    hCompletionPort = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, NULL, 0);
    if (hCompletionPort == NULL) {
        MyErrExit("CreateIoCompletionPort");
    }

    // Create completion port worker threads.
    //
    CreateWorkers(hCompletionPort);
    CreatePipeConnection(hCompletionPort, 0);

    printf("TRACE: Ready for clients.  Press Ctrl-C to stop.\n");

    /////////////////////////////////////////////////////////// Validate DLLs.
    //
    CHAR szTmpPath[MAX_PATH];
    CHAR szExePath[MAX_PATH];
    CHAR szDllPath[MAX_PATH];
    PCHAR pszFilePart = NULL;

    if (!GetModuleFileName(NULL, szTmpPath, ARRAYOF(szTmpPath))) {
        printf("TRACE: Couldn't retreive exe name.\n");
        return 9002;
    }
    if (!GetFullPathName(szTmpPath, ARRAYOF(szExePath), szExePath, &pszFilePart) ||
        pszFilePart == NULL) {
        printf("TRACE: Error: %s is not a valid path name..\n", szTmpPath);
        return 9002;
    }

#ifdef _CRT_INSECURE_DEPRECATE
    strcpy_s(pszFilePart, szExePath + ARRAYOF(szExePath) - pszFilePart, "Trace.Detours.dll");
    strcpy_s(szDllPath, ARRAYOF(szDllPath), szExePath);
#else
    strcpy(pszFilePart, "Trace.Detours.dll");
    strcpy(szDllPath, szExePath);
#endif

    //////////////////////////////////////////////////////////////////////////
    STARTUPINFO si;
    PROCESS_INFORMATION pi;
    CHAR szCommand[2048];
    CHAR szExe[MAX_PATH];
    CHAR szFullExe[MAX_PATH] = "\0";
    PCHAR pszFileExe = NULL;

    ZeroMemory(&si, sizeof(si));
    ZeroMemory(&pi, sizeof(pi));
    si.cb = sizeof(si);

    szCommand[0] = L'\0';

#ifdef _CRT_INSECURE_DEPRECATE
    strcpy_s(szExe, sizeof(szExe), argv[arg]);
#else
    strcpy(szExe, argv[arg]);
#endif
    for (; arg < argc; arg++) {
        if (strchr(argv[arg], ' ') != NULL || strchr(argv[arg], '\t') != NULL) {
#ifdef _CRT_INSECURE_DEPRECATE
            strcat_s(szCommand, sizeof(szCommand), "\"");
            strcat_s(szCommand, sizeof(szCommand), argv[arg]);
            strcat_s(szCommand, sizeof(szCommand), "\"");
#else
            strcat(szCommand, "\"");
            strcat(szCommand, argv[arg]);
            strcat(szCommand, "\"");
#endif
        }
        else {
#ifdef _CRT_INSECURE_DEPRECATE
            strcat_s(szCommand, sizeof(szCommand), argv[arg]);
#else
            strcat(szCommand, argv[arg]);
#endif
        }

        if (arg + 1 < argc) {
#ifdef _CRT_INSECURE_DEPRECATE
            strcat_s(szCommand, sizeof(szCommand), " ");
#else
            strcat(szCommand, " ");
#endif
        }
    }
    printf("TRACE: Starting: `%s'\n", szCommand);
    printf("TRACE:   with `%s'\n", szDllPath);
    fflush(stdout);

    DWORD dwFlags = CREATE_DEFAULT_ERROR_MODE | CREATE_SUSPENDED;

    SetLastError(0);
    SearchPath(NULL, szExe, ".exe", ARRAYOF(szFullExe), szFullExe, &pszFileExe);


    if (!DetourCreateProcessWithDll(szFullExe[0] ? szFullExe : NULL, szCommand,
                                    NULL, NULL, TRUE, dwFlags, NULL, NULL,
                                    &si, &pi, NULL, szDllPath, NULL)) {
        printf("TRACE: DetourCreateProcessWithDll failed: %d\n", GetLastError());
        ExitProcess(9007);
    }

    ZeroMemory(&s_Payload, sizeof(s_Payload));
    s_Payload.nParentProcessId = GetCurrentProcessId();
    s_Payload.nTraceProcessId = GetCurrentProcessId();
    s_Payload.nGeneology = 1;
    s_Payload.rGeneology[0] = 0;
    wcscpy_s(s_Payload.wzStdin, ARRAYOF(s_Payload.wzStdin), L"\\\\.\\CONIN$");
    wcscpy_s(s_Payload.wzStdout, ARRAYOF(s_Payload.wzStdout), L"\\\\.\\CONOUT$");
    wcscpy_s(s_Payload.wzStderr, ARRAYOF(s_Payload.wzStderr), L"\\\\.\\CONOUT$");
    wcscpy_s(s_Payload.wzParents, ARRAYOF(s_Payload.wzParents), L"");
    CopyEnvironment(s_Payload.wzzDrop, wzzDrop);
    LPWCH pwStrings = GetEnvironmentStringsW();
    CopyEnvironment(s_Payload.wzzEnvironment, pwStrings);
    FreeEnvironmentStringsW(pwStrings);

    if (!DetourCopyPayloadToProcess(pi.hProcess, s_guidTrace,
                                    &s_Payload, sizeof(s_Payload))) {
        printf("TRACE: DetourCopyPayloadToProcess failed: %d\n", GetLastError());
        ExitProcess(9008);
    }

    ResumeThread(pi.hThread);

    WaitForSingleObject(pi.hProcess, INFINITE);

    DWORD dwResult = 0;
    if (!GetExitCodeProcess(pi.hProcess, &dwResult)) {
        printf("TRACE: GetExitCodeProcess failed: %d\n", GetLastError());
        return 9008;
    }

    printf("TRACE: %d processes.\n", s_nTotalClients);
    printf("TRACE %s.xml\r\n", s_szLogFile);
    
    CHAR finalLogFile[MAX_PATH];
#ifdef _CRT_INSECURE_DEPRECATE
    sprintf_s(finalLogFile, ARRAYOF(finalLogFile), "%s.xml", s_szLogFile);
#else
    sprintf(finalLogFile, "%s.xml", s_szLogFile);
#endif

    HANDLE hFinal = CreateFile(finalLogFile,GENERIC_WRITE,FILE_SHARE_READ,NULL,CREATE_ALWAYS,FILE_ATTRIBUTE_NORMAL |FILE_FLAG_SEQUENTIAL_SCAN,NULL);    
    
    CHAR eachLogFile[MAX_PATH];
    WIN32_FILE_ATTRIBUTE_DATA fileData;
    BOOL rc;
    DWORD fileSize;
    int i;
    WriteOut( hFinal, "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n<Traces>\n" );
    for( i = 0; i < s_nTotalClients;i++ ){ 
        
        
#ifdef _CRT_INSECURE_DEPRECATE
    sprintf_s(eachLogFile, ARRAYOF(eachLogFile), "%s\\%s.%08d.xml", s_szTempPath, s_szJustFileName, i);
#else
    sprintf(eachLogFile, "%s\\%s.%08d.xml", s_szTempPath, s_szJustFileName, i);
#endif
        rc = GetFileAttributesEx(eachLogFile, GetFileExInfoStandard, &fileData );
        fileSize = fileData.nFileSizeLow;
        char* buffer = new char[fileSize*4];
        HANDLE hLog = CreateFile(eachLogFile,GENERIC_READ,FILE_SHARE_READ,NULL,OPEN_EXISTING,FILE_ATTRIBUTE_NORMAL,NULL);    
        ReadFile( hLog , buffer, fileSize, &fileSize, NULL);
        CloseHandle(hLog);
        DeleteFile( eachLogFile );
        WriteFile(hFinal, buffer, fileSize, &fileSize, NULL);
        delete[] buffer;
    }
    
    WriteOut( hFinal, "</Traces>" );
    CloseHandle(hFinal);
    return dwResult;
}
